from torch import nn

class MLPModel(nn.Module):
    def __init__(self, n = 2):
        super(MLPModel, self).__init__()
        self.fc1 = nn.Linear(784, 256)
        self.conv1 = nn.Conv2d(1,16,3,1,1)
        self.conv2 = nn.Conv2d(16,16,3,1,1)
        self.conv3 = nn.Conv2d(16,1,3,1,1)
        self.n_layer = n
        self.relu1 = nn.ReLU()
        self.fc2 = nn.Linear(256, 10)

    def forward(self, x):
        x = self.fc1(x)
        x = self.conv1(x.reshape(1,16,16))
        for i in range(self.n_layer):
            x = self.conv2(x)
        x = self.conv3(x).flatten().reshape(1,-1)
        x = self.relu1(x)
        x = self.fc2(x)
        return x
